import torch
from torch import optim
from tqdm import tqdm
from src import *

device='cuda' if torch.cuda.is_available() else 'cpu'
model=NaiveNetwork(n_chan=3).to(device)
model2=NaiveNetwork(n_chan=3).to(device)

# ---------- global parameters ----------
# TODO: implement load except for fix
max_epoch = 3000     # training epochs
lr = 0.001           # learning rate
step_size = 1000     # number of epochs at which learning rate decays
gamma = 0.5          # factor by which learning rate decays
save_path=r'save.png'


optimizer = optim.Adam(model.parameters(), lr=lr)
optimizer2 = optim.Adam(model2.parameters(), lr=lr)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=step_size, gamma=gamma)
scheduler2 = optim.lr_scheduler.StepLR(optimizer2, step_size=step_size, gamma=gamma)

Picture2Tensor(img_path=r'data/raw/牛角蛋糕.png', pt_path=r'data/cache/牛角蛋糕.pt')
clean_image:torch.tensor=torch.load(r'data/cache/牛角蛋糕.pt').to(device)
if len(clean_image.shape)==3:
    clean_image=clean_image.unsqueeze(0)
noisy_image:torch.tensor=add_noise(clean_image, noise_type='gauss', noise_level=50)

def train(model, optimizer, noisy_img):
    
  loss = loss_func(noisy_img, model)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

  return loss.item()

def train_with_var(model, optimizer, noisy_img):
    
  loss = loss_func_with_var(noisy_img, model)
  optimizer.zero_grad()
  loss.backward()
  optimizer.step()

  return loss.item()

# ---------- 执行训练 ----------
noisy_image=noisy_image


for epoch in tqdm(range(max_epoch)):
    train(model, optimizer, noisy_image)
    scheduler.step()

    
# ---------- 获取干净的图片 ----------
denoised_img=noisy_image-model(noisy_image)
denoised_img=denoised_img

# ---------- 二次denoise ----------
noisy_image2=denoised_img
noisy_image2=noisy_image2.detach()
for epoch in tqdm(range(max_epoch)):
    train_with_var(model2, optimizer2, noisy_image2)
    scheduler2.step()

denoised_img2=noisy_image2-model2(noisy_image2)
denoised_img2=denoised_img2

# ---------- 获得PSNR分数 ----------

psnr_clean=getPSNR(
    noisy_img=clean_image,
    clean_img=clean_image
)
psnr_noise=getPSNR(
    noisy_img=noisy_image,
    clean_img=clean_image
)
psnr_denoised=getPSNR(
    noisy_img=denoised_img,
    clean_img=clean_image
)
psnr_denoised2=getPSNR(
    noisy_img=denoised_img2,
    clean_img=clean_image
)

# ---------- 保存结果 ----------
dc=[
    {
        'img':Tensor2Picture(clean_image),
        'label':psnr_clean
    },
    {
        'img':Tensor2Picture(noisy_image),
        'label':psnr_noise
    },
    {
        'img':Tensor2Picture(denoised_img),
        'label':psnr_denoised
    },
    {
        'img':Tensor2Picture(denoised_img2),
        'label':psnr_denoised2
    }
]
plot(dc, output_path=save_path)
